package org.jeecg.common.es;

import org.apache.lucene.queryparser.flexible.core.util.StringUtils;
import org.apache.lucene.queryparser.flexible.standard.QueryParserUtil;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.springframework.data.elasticsearch.core.query.Criteria;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

import java.util.*;

import static org.elasticsearch.index.query.Operator.AND;
import static org.elasticsearch.index.query.QueryBuilders.*;

/**
 * @author wangzun1.vendor@sensetime.com
 * @Description:
 * @date 2022/8/10
 */
public class CriteriaQueryProcessor {
    public QueryBuilder createQueryFromCriteria(Criteria criteria) {
        Assert.notNull(criteria, "criteria must not be null");
        List<QueryBuilder> shouldQueryBuilderList = new LinkedList<>();
        List<QueryBuilder> mustNotQueryBuilderList = new LinkedList<>();
        List<QueryBuilder> mustQueryBuilderList = new LinkedList<>();
        ListIterator<Criteria> chainIterator = criteria.getCriteriaChain().listIterator();
        QueryBuilder firstQuery = null;
        boolean negateFirstQuery = false;

        while (chainIterator.hasNext()) {
            Criteria chainedCriteria = chainIterator.next();
            QueryBuilder queryFragmentForCriteria = createQueryFragmentForCriteria(chainedCriteria);
            if (queryFragmentForCriteria != null) {
                if (firstQuery == null) {
                    firstQuery = queryFragmentForCriteria;
                    negateFirstQuery = chainedCriteria.isNegating();
                    continue;
                }
                if (chainedCriteria.isOr()) {
                    shouldQueryBuilderList.add(queryFragmentForCriteria);
                } else if (chainedCriteria.isNegating()) {
                    mustNotQueryBuilderList.add(queryFragmentForCriteria);
                } else {
                    mustQueryBuilderList.add(queryFragmentForCriteria);
                }
            }
        }

        if (firstQuery != null) {
            if (!shouldQueryBuilderList.isEmpty() && mustNotQueryBuilderList.isEmpty()
                && mustQueryBuilderList.isEmpty()) {
                shouldQueryBuilderList.add(0, firstQuery);
            } else {
                if (negateFirstQuery) {
                    mustNotQueryBuilderList.add(0, firstQuery);
                } else {
                    mustQueryBuilderList.add(0, firstQuery);
                }
            }
        }

        BoolQueryBuilder query = null;

        if (!shouldQueryBuilderList.isEmpty() || !mustNotQueryBuilderList.isEmpty()
            || !mustQueryBuilderList.isEmpty()) {

            query = boolQuery();

            for (QueryBuilder qb : shouldQueryBuilderList) {
                query.should(qb);
            }
            for (QueryBuilder qb : mustNotQueryBuilderList) {
                query.mustNot(qb);
            }
            for (QueryBuilder qb : mustQueryBuilderList) {
                query.must(qb);
            }
        }

        return query;
    }

    @Nullable
    private QueryBuilder createQueryFragmentForCriteria(Criteria chainedCriteria) {
        if (chainedCriteria.getQueryCriteriaEntries().isEmpty())
            return null;

        Iterator<Criteria.CriteriaEntry> it = chainedCriteria.getQueryCriteriaEntries().iterator();
        boolean singeEntryCriteria = (chainedCriteria.getQueryCriteriaEntries().size() == 1);

        String fieldName = chainedCriteria.getField().getName();
        Assert.notNull(fieldName, "Unknown field");
        QueryBuilder query = null;

        if (singeEntryCriteria) {
            Criteria.CriteriaEntry entry = it.next();
            query = processCriteriaEntry(entry, fieldName);
        } else {
            query = boolQuery();
            while (it.hasNext()) {
                Criteria.CriteriaEntry entry = it.next();
                ((BoolQueryBuilder)query).must(processCriteriaEntry(entry, fieldName));
            }
        }

        addBoost(query, chainedCriteria.getBoost());
        return query;
    }

    @Nullable
    private QueryBuilder processCriteriaEntry(Criteria.CriteriaEntry entry, String fieldName) {
        Criteria.OperationKey key = entry.getKey();
        Object value = entry.getValue();

        if (value == null) {

            if (key == Criteria.OperationKey.EXISTS) {
                return existsQuery(fieldName);
            } else {
                return null;
            }
        }

        String searchText = QueryParserUtil.escape(value.toString());

        QueryBuilder query = null;

        switch (key) {
            case EQUALS:
                query = queryStringQuery(searchText).field(fieldName).defaultOperator(AND);
                break;
            case CONTAINS:
                query = queryStringQuery('*' + searchText + '*').field(fieldName).analyzeWildcard(true);
                break;
            case STARTS_WITH:
                query = queryStringQuery(searchText + '*').field(fieldName).analyzeWildcard(true);
                break;
            case ENDS_WITH:
                query = queryStringQuery('*' + searchText).field(fieldName).analyzeWildcard(true);
                break;
            case EXPRESSION:
                query = queryStringQuery(value.toString()).field(fieldName);
                break;
            case LESS_EQUAL:
                query = rangeQuery(fieldName).lte(value);
                break;
            case GREATER_EQUAL:
                query = rangeQuery(fieldName).gte(value);
                break;
            case BETWEEN:
                Object[] ranges = (Object[])value;
                query = rangeQuery(fieldName).from(ranges[0]).to(ranges[1]);
                break;
            case LESS:
                query = rangeQuery(fieldName).lt(value);
                break;
            case GREATER:
                query = rangeQuery(fieldName).gt(value);
                break;
            case FUZZY:
                query = fuzzyQuery(fieldName, searchText);
                break;
            case IN:
                query = boolQuery().must(termsQuery(fieldName, toStringList((Iterable<Object>)value)));
                break;
            case NOT_IN:
                query = boolQuery().mustNot(termsQuery(fieldName, toStringList((Iterable<Object>)value)));
                break;
        }
        return query;
    }

    private static List<String> toStringList(Iterable<?> iterable) {
        List<String> list = new ArrayList<>();
        for (Object item : iterable) {
            list.add(StringUtils.toString(item));
        }
        return list;
    }

    private void addBoost(QueryBuilder query, float boost) {
        if (Float.isNaN(boost)) {
            return;
        }
        query.boost(boost);
    }
}
